#!/usr/bin/env python3
"""
数据集集成脚本
集成HuggingFace简历数据集和其他数据源
"""

import asyncio
import json
import os
from typing import Dict, Any, List
from pathlib import Path

from src.core.config import config
from src.agents.agent_manager import agent_manager
from src.lightrag_engine.ingestion.pipeline import IngestionPipeline

class DatasetIntegrator:
    """数据集集成器"""
    
    def __init__(self):
        self.agent_manager = agent_manager
        self.ingestion_pipeline = None
    
    async def initialize(self):
        """初始化系统"""
        print("🚀 Initializing BoLe HR Platform for dataset integration...")
        await self.agent_manager.initialize()
        
        orchestrator = self.agent_manager.get_orchestrator()
        if orchestrator and orchestrator.rag_controller:
            self.ingestion_pipeline = orchestrator.rag_controller.ingestion_pipeline
        else:
            raise RuntimeError("Failed to initialize ingestion pipeline")
    
    async def integrate_huggingface_dataset(self, dataset_name: str = "opensporks/resumes"):
        """集成HuggingFace简历数据集"""
        print(f"📊 Integrating HuggingFace dataset: {dataset_name}")
        
        try:
            # 使用ingestion pipeline处理HuggingFace数据集
            await self.ingestion_pipeline.process_huggingface_dataset(
                dataset_name=dataset_name,
                metadata={
                    "source": "huggingface",
                    "dataset": dataset_name,
                    "integration_date": "2025-01-01",
                    "description": "Resume dataset from HuggingFace"
                }
            )
            
            print(f"✅ Successfully integrated {dataset_name}")
            
            # 获取统计信息
            stats = await self.ingestion_pipeline.get_pipeline_stats()
            print(f"📈 Pipeline stats: {json.dumps(stats, indent=2)}")
            
        except Exception as e:
            print(f"❌ Failed to integrate HuggingFace dataset: {str(e)}")
            raise
    
    async def create_sample_resumes(self, output_dir: str = "data/samples"):
        """创建示例简历数据"""
        print(f"📝 Creating sample resume data in {output_dir}")
        
        # 确保输出目录存在
        Path(output_dir).mkdir(parents=True, exist_ok=True)
        
        sample_resumes = [
            {
                "id": "resume_001",
                "name": "张三 (Zhang San)",
                "position": "Senior Python Developer",
                "content": """
张三 - 高级Python开发工程师

联系信息：
- Email: zhangsan@email.com  
- Phone: +86 138-0000-0000
- Location: 北京市朝阳区

专业技能：
- 编程语言：Python, JavaScript, Java, Go
- 框架技术：Django, FastAPI, React, Vue.js
- 数据库：MySQL, PostgreSQL, MongoDB, Redis
- 云平台：AWS, Azure, 阿里云
- 机器学习：scikit-learn, TensorFlow, PyTorch
- 开发工具：Git, Docker, Kubernetes

工作经验：
2021-2024 | 字节跳动 | 高级后端开发工程师
- 负责推荐系统后端开发，处理日活过亿的用户请求
- 使用Python+Django构建微服务架构，提升系统吞吐量30%
- 优化数据库查询性能，减少响应时间50%
- 带领5人团队完成多个核心项目

2019-2021 | 美团 | Python开发工程师  
- 开发外卖配送优化算法，提升配送效率15%
- 使用Flask构建RESTful API，支撑千万级并发
- 实现实时数据处理管道，处理每日TB级数据

教育背景：
2015-2019 | 清华大学 | 计算机科学与技术 | 学士学位
- GPA: 3.8/4.0
- 核心课程：数据结构、算法设计、机器学习、分布式系统

项目经历：
- 智能推荐系统：基于深度学习的个性化推荐，提升点击率25%
- 实时数据分析平台：处理实时流数据，支持复杂查询和可视化
- 微服务架构重构：将单体应用拆分为微服务，提升系统可维护性

语言能力：
- 中文：母语
- 英语：CET-6，能够流利进行技术交流

证书：
- AWS Certified Solutions Architect
- Oracle Certified Professional Java Programmer
                """,
                "skills": ["Python", "Django", "FastAPI", "机器学习", "微服务", "AWS"],
                "experience_years": 5,
                "education": "清华大学计算机科学学士",
                "location": "北京"
            },
            {
                "id": "resume_002", 
                "name": "李四 (Li Si)",
                "position": "Full Stack Developer",
                "content": """
李四 - 全栈开发工程师

个人信息：
- Email: lisi@email.com
- Phone: +86 139-0000-0000  
- GitHub: github.com/lisi
- Location: 上海市浦东新区

技术栈：
前端：React, Vue.js, TypeScript, HTML5, CSS3, Webpack
后端：Node.js, Express, Koa, Python, Flask
数据库：MySQL, MongoDB, Redis
运维：Docker, Kubernetes, CI/CD, Linux
其他：Git, Agile, TDD

工作经历：
2020-2024 | 腾讯 | 全栈开发工程师
- 负责QQ音乐Web端和移动端开发，用户量过5亿
- 使用React+TypeScript重构前端架构，提升开发效率40% 
- 开发Node.js后端服务，支持高并发音乐流媒体服务
- 实现前后端分离架构，优化页面加载速度60%

2018-2020 | 滴滴出行 | 前端开发工程师
- 开发司机端H5应用，覆盖全国300+城市
- 使用Vue.js构建响应式界面，适配多种设备
- 优化前端性能，首屏加载时间减少40%
- 参与移动端原生应用开发，熟悉React Native

教育背景：
2014-2018 | 上海交通大学 | 软件工程 | 学士学位
- 主修课程：Web开发、移动应用开发、软件架构
- 毕业设计：基于微信小程序的校园生活服务平台

个人项目：
- 在线教育平台：React+Node.js全栈开发，支持视频直播
- 个人博客系统：Vue.js+Express，包含后台管理系统
- 开源组件库：发布NPM包，GitHub Star 500+

语言：
- 中文：母语  
- 英语：CET-6，良好的读写能力

兴趣爱好：
- 开源贡献者，活跃于GitHub社区
- 技术博客写作，CSDN博客专家
- 摄影、旅行、读书
                """,
                "skills": ["React", "Vue.js", "Node.js", "TypeScript", "全栈开发"],
                "experience_years": 6,
                "education": "上海交通大学软件工程学士", 
                "location": "上海"
            },
            {
                "id": "resume_003",
                "name": "王五 (Wang Wu)", 
                "position": "Data Scientist",
                "content": """
王五 - 数据科学家

联系方式：
- Email: wangwu@email.com
- LinkedIn: linkedin.com/in/wangwu
- Phone: +86 137-0000-0000
- Location: 深圳市南山区

核心技能：
- 编程语言：Python, R, SQL, Scala
- 机器学习：scikit-learn, TensorFlow, PyTorch, XGBoost
- 数据处理：Pandas, NumPy, Spark, Hadoop
- 可视化：Matplotlib, Seaborn, Plotly, Tableau
- 云平台：AWS, GCP, Azure
- 统计分析：假设检验, A/B测试, 时间序列分析

工作经验：
2022-2024 | 华为 | 高级数据科学家
- 领导AI算法团队，开发智能客服NLP模型，准确率达95%
- 构建用户行为预测模型，提升产品转化率30%
- 使用深度学习优化推荐算法，增加用户停留时间25%
- 建立实时数据监控体系，及时发现业务异常

2020-2022 | 平安科技 | 数据科学家
- 开发风控模型，降低信贷违约率20%
- 使用机器学习进行保险定价，提升保费收入15%
- 构建客户画像系统，支持精准营销
- 参与数据平台建设，处理PB级数据

2018-2020 | 京东 | 数据分析师
- 负责电商推荐系统数据分析，提升GMV 10%
- 进行用户行为分析，优化产品功能设计
- 开发销售预测模型，提高库存周转率
- 使用A/B测试验证产品改进效果

教育背景：
2014-2018 | 北京大学 | 统计学 | 学士学位
- 专业课程：概率论、数理统计、回归分析、时间序列
- 毕业论文：基于深度学习的股票价格预测研究

2018-2020 | 斯坦福大学 | 数据科学 | 硕士学位  
- GPA: 3.9/4.0
- 研究方向：机器学习、自然语言处理
- 导师：Andrew Ng教授

项目经历：
- 智能客服系统：基于BERT的意图识别和情感分析
- 金融风控平台：集成多种机器学习算法的风险评估系统
- 实时推荐引擎：使用协同过滤和深度学习的混合推荐
- 数据可视化大屏：实时展示业务关键指标

发表论文：
- "Deep Learning for Customer Behavior Prediction" - KDD 2023
- "Real-time Fraud Detection Using Machine Learning" - ICML 2022

证书与荣誉：
- Google Cloud Professional Data Engineer
- AWS Certified Machine Learning Specialty  
- 华为技术创新奖 (2023)
- 平安科技优秀员工 (2021)

语言能力：
- 中文：母语
- 英语：托福110分，流利的听说读写能力
                """,
                "skills": ["Python", "机器学习", "深度学习", "数据分析", "NLP", "TensorFlow"],
                "experience_years": 6,
                "education": "北京大学统计学学士，斯坦福大学数据科学硕士",
                "location": "深圳"
            }
        ]
        
        # 保存样本简历
        for resume in sample_resumes:
            file_path = os.path.join(output_dir, f"{resume['id']}.json")
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(resume, f, ensure_ascii=False, indent=2)
        
        print(f"✅ Created {len(sample_resumes)} sample resumes")
        return sample_resumes
    
    async def create_job_requirements(self, output_dir: str = "data/samples"):
        """创建职位需求样例"""
        print(f"💼 Creating job requirement samples in {output_dir}")
        
        Path(output_dir).mkdir(parents=True, exist_ok=True)
        
        job_requirements = [
            {
                "id": "job_001",
                "title": "高级Python后端开发工程师",
                "company": "字节跳动",
                "location": "北京",
                "salary_range": "35-50K",
                "requirements": """
职位职责：
1. 负责后端服务架构设计和核心功能开发
2. 参与系统性能优化和稳定性改进
3. 指导初级开发人员技术成长
4. 参与技术方案评审和代码Review

任职要求：
1. 本科以上学历，计算机相关专业
2. 5年以上Python开发经验，熟悉Django/Flask/FastAPI
3. 熟悉MySQL、Redis、MongoDB等数据库
4. 熟悉微服务架构，有大规模分布式系统经验
5. 熟悉云平台（AWS/阿里云），了解容器化部署
6. 具备较强的问题分析和解决能力
7. 有团队管理经验者优先

技能要求：
- 必须：Python, Django/Flask, MySQL, Redis, Git
- 优先：Kubernetes, Docker, 微服务, 机器学习
- 加分：Go, Java, 大数据处理经验
                """,
                "key_skills": ["Python", "Django", "微服务", "MySQL", "Redis", "云平台"],
                "experience_required": "5+",
                "education_required": "本科以上",
                "team_size": "10-20人"
            },
            {
                "id": "job_002", 
                "title": "全栈开发工程师",
                "company": "腾讯",
                "location": "深圳",
                "salary_range": "25-40K",
                "requirements": """
岗位描述：
负责公司核心产品的前后端开发，参与产品架构设计和技术方案制定。

工作内容：
1. 负责Web前端和移动端H5开发
2. 开发Node.js后端API服务
3. 参与产品需求分析和技术方案设计
4. 进行代码重构和性能优化
5. 配合产品经理和设计师完成项目开发

岗位要求：
1. 本科及以上学历，3年以上全栈开发经验
2. 精通JavaScript, HTML5, CSS3
3. 熟练使用React/Vue.js等前端框架
4. 熟悉Node.js后端开发，了解Express/Koa
5. 熟悉MySQL/MongoDB数据库操作
6. 了解RESTful API设计原则
7. 有移动端开发经验者优先
8. 良好的团队协作和沟通能力

技术栈：
前端：React, Vue.js, TypeScript, Webpack
后端：Node.js, Express, Python
数据库：MySQL, MongoDB, Redis
工具：Git, Docker, Linux
                """,
                "key_skills": ["JavaScript", "React", "Vue.js", "Node.js", "全栈开发"],
                "experience_required": "3+", 
                "education_required": "本科以上",
                "team_size": "5-15人"
            },
            {
                "id": "job_003",
                "title": "数据科学家",
                "company": "华为",
                "location": "上海",
                "salary_range": "40-60K", 
                "requirements": """
职位描述：
负责公司AI产品的算法研发，参与机器学习模型的设计、开发和优化。

主要职责：
1. 进行数据挖掘和特征工程
2. 设计和实现机器学习算法模型
3. 进行A/B测试和模型效果评估
4. 参与AI产品的技术方案设计
5. 跟踪前沿技术，推动算法创新

任职要求：
1. 硕士以上学历，统计学/数学/计算机相关专业
2. 3年以上数据科学或机器学习工作经验
3. 精通Python/R，熟悉pandas, numpy, scikit-learn
4. 熟悉深度学习框架TensorFlow/PyTorch
5. 具备扎实的统计学和机器学习理论基础
6. 有大规模数据处理经验，熟悉Spark/Hadoop
7. 有NLP或计算机视觉经验者优先
8. 良好的英文文献阅读能力

加分项：
- 发表过顶级会议/期刊论文
- 参加过Kaggle等数据竞赛并获奖
- 有推荐系统或搜索算法经验
- 熟悉云平台机器学习服务
                """,
                "key_skills": ["Python", "机器学习", "深度学习", "TensorFlow", "数据挖掘"],
                "experience_required": "3+",
                "education_required": "硕士以上", 
                "team_size": "3-8人"
            }
        ]
        
        # 保存职位需求
        for job in job_requirements:
            file_path = os.path.join(output_dir, f"{job['id']}.json")
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(job, f, ensure_ascii=False, indent=2)
        
        print(f"✅ Created {len(job_requirements)} job requirement samples")
        return job_requirements
    
    async def test_analysis_pipeline(self, resume_samples: List[Dict], job_samples: List[Dict]):
        """测试分析管道"""
        print("🧪 Testing analysis pipeline...")
        
        orchestrator = self.agent_manager.get_orchestrator()
        if not orchestrator:
            raise RuntimeError("Orchestrator not available")
        
        # 测试简历分析
        for resume in resume_samples[:2]:  # 测试前两个简历
            print(f"📄 Testing resume analysis for: {resume['name']}")
            
            from src.core.interfaces import Task
            task = Task(
                id=f"test_resume_{resume['id']}",
                description=resume['content'],
                task_type="resume_analysis",
                metadata={
                    "candidate_name": resume['name'],
                    "position": resume['position']
                }
            )
            
            try:
                result = await orchestrator.orchestrate_task(task)
                print(f"✅ Resume analysis completed: {result.status}")
                if result.result:
                    print(f"   Analysis summary: {str(result.result)[:200]}...")
            except Exception as e:
                print(f"❌ Resume analysis failed: {str(e)}")
        
        # 测试候选人匹配
        print("\n🎯 Testing candidate-job matching...")
        try:
            match_task = Task(
                id="test_matching",
                description=f"Match candidate {resume_samples[0]['name']} with job {job_samples[0]['title']}",
                task_type="candidate_matching",
                metadata={
                    "candidate_info": resume_samples[0]['content'],
                    "job_requirements": job_samples[0]['requirements']
                }
            )
            
            result = await orchestrator.orchestrate_task(match_task)
            print(f"✅ Candidate matching completed: {result.status}")
            
        except Exception as e:
            print(f"❌ Candidate matching failed: {str(e)}")
    
    async def cleanup(self):
        """清理资源"""
        print("🧹 Cleaning up resources...")
        await self.agent_manager.shutdown()

async def main():
    """主函数"""
    integrator = DatasetIntegrator()
    
    try:
        # 初始化系统
        await integrator.initialize()
        
        # 创建样本数据
        resume_samples = await integrator.create_sample_resumes()
        job_samples = await integrator.create_job_requirements()
        
        # 集成HuggingFace数据集（可选，需要网络连接）
        try:
            await integrator.integrate_huggingface_dataset()
        except Exception as e:
            print(f"⚠️  HuggingFace integration skipped: {str(e)}")
        
        # 测试分析管道
        await integrator.test_analysis_pipeline(resume_samples, job_samples)
        
        print("\n🎉 Dataset integration and testing completed successfully!")
        
    except Exception as e:
        print(f"❌ Integration failed: {str(e)}")
        raise
        
    finally:
        await integrator.cleanup()

if __name__ == "__main__":
    asyncio.run(main())